#include <cstdint>
#include <iostream>
#include <fstream>
#include <string>
#include <cstring>
#include <cerrno>
#include <byteswap.h>
#include <sstream>
#include <endian.h>
#include <algorithm>
#include "rpihw.h"

#define LINE_WIDTH_MAX 80
#define HW_VER_STRING "Revision"

#define PERIPH_BASE_RPI 0x20000000
#define PERIPH_BASE_RPI2 0x3f000000
#define PERIPH_BASE_RPI4 0xfe000000

#define VIDEOCORE_BASE_RPI 0x40000000
#define VIDEOCORE_BASE_RPI2 0xc0000000

#define RPI_MANUFACTURER_MASK (0xf << 16)
#define RPI_WARRANTY_MASK (0x3 << 24)


static const rpi_hw_t rpi_hw_info[] = {
    //
    // Raspberry Pi 400
    //
    {0xc03130, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Pi 400 - 4GB v1.0"},
    {0xc03131, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Pi 400 - 4GB v1.1"},
    //
    // Raspberry Pi 4
    //
    {0xA03111, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Pi 4 Model B - 1GB v1.1"},
    {0xB03111, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Pi 4 Model B - 2GB v.1.1"},
    {0xC03111, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Pi 4 Model B - 4GB v1.1"},
    {0xA03112, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Pi 4 Model B - 1GB v1.2"},
    {0xB03112, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Pi 4 Model B - 2GB v.1.2"},
    {0xC03112, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Pi 4 Model B - 4GB v1.2"},
    {0xb03114, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Pi 4 Model B - 2GB v1.4"},
    {0xD03114, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Pi 4 Model B - 8GB v1.4"},
    {0xc03114, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Pi 4 Model B - 4GB v1.4"},
    {0xa03115, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Pi 4 Model B - 1GB v1.5"},
    {0xb03115, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Pi 4 Model B - 2GB v1.5"},
    {0xc03115, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Pi 4 Model B - 4GB v1.5"},
    {0xd03115, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Pi 4 Model B - 8GB v1.5"},

    //
    // Compute Module 4
    //
    {0xa03140, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Compute Module 4 v1.0 eMMC"},
    {0xb03140, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Compute Module 4 v1.0 Lite"},
    {0xc03140, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Compute Module 4 v1.0 WiFi"},
    {0xd03140, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Compute Module 4 v1.0 WiFi 8GB"},
    {0xa03141, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Compute Module 4 Rev 1.1"},
    {0xb03141, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Compute Module 4 Rev 1.1"},
    {0xc03141, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Compute Module 4 Rev 1.1"},
    {0xd03141, RPI_HWVER_TYPE_PI4, PERIPH_BASE_RPI4, VIDEOCORE_BASE_RPI2, "Compute Module 4 Rev 1.1"},

    //
    // Model B Rev 1.0
    //
    {0x02, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model B"},
    {0x03, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model B"},

    //
    // Model B Rev 2.0
    //
    {0x04, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model B"},
    {0x05, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model B"},
    {0x06, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model B"},

    //
    // Model A
    //
    {0x07, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model A"},
    {0x08, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model A"},
    {0x09, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model A"},

    //
    // Model B
    //
    {0x0d, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model B"},
    {0x0e, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model B"},
    {0x0f, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model B"},

    //
    // Model B+
    //
    {0x10, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model B+"},
    {0x13, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model B+"},
    {0x900032, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model B+"},
    
    //
    // Compute Module
    //
    {0x11, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Compute Module 1"},
    {0x14, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Compute Module 1"},
    {0xa020a0, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Compute Module 3/L3"},
    {0xa02100, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Compute Module 3+"},
    {0xa220a0, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Compute Module 3 Rev 1.0"},

    //
    // Pi Zero
    //
    {0x900092, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Pi Zero v1.2"},
    {0x900093, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Pi Zero v1.3"},
    {0x920093, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Pi Zero v1.3"},
    {0x9200c1, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Pi Zero W v1.1"},
    {0x9000c1, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Pi Zero W v1.1"},
    //
    // Model Zero 2 W
    //
    {0x902120, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Pi Zero 2 W v1.0"},

    //
    // Model A+
    //
    {0x12, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model A+"},
    {0x15, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model A+"},
    {0x900021, RPI_HWVER_TYPE_PI1, PERIPH_BASE_RPI, VIDEOCORE_BASE_RPI, "Model A+"},

    //
    // Pi 2 Model B
    //
    {0xa01041, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Pi 2"},
    {0xa01040, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Pi 2"},
    {0xa21041, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Pi 2"},

    //
    // Pi 2 with BCM2837
    //
    {0xa22042, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Pi 2"},

    //
    // Pi 3 Model B
    //
    {0xa020d4, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Pi 3 B+"},
    {0xa020d3, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Pi 3 B+"},
    {0xa32082, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Pi 3 B"},
    {0xa02082, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Pi 3"},
    {0xa02083, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Pi 3"},
    {0xa22082, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Pi 3"},
    {0xa22083, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Pi 3"},
    {0x9020e0, RPI_HWVER_TYPE_PI2, PERIPH_BASE_RPI2, VIDEOCORE_BASE_RPI2, "Model 3 A+"}
};


// 获取 rpi_hw_info 的大小
constexpr size_t rpi_hw_info_size = sizeof(rpi_hw_info) / sizeof(rpi_hw_info[0]);

const rpi_hw_t* rpi_hw_detect() {
    const rpi_hw_t* result = nullptr;
    uint32_t rev = 0;
    std::ifstream file;
    std::string line;

#ifdef __aarch64__
    file.open("/proc/device-tree/system/linux,revision", std::ios::binary);
    if (!file.is_open()) {
        return nullptr;
    }
    
    file.read(reinterpret_cast<char*>(&rev), sizeof(uint32_t));
    if (file.gcount() != sizeof(uint32_t)) {
        file.close();
        return nullptr;
    }
    
    #if __BYTE_ORDER == __LITTLE_ENDIAN
        rev = __bswap_32(rev); // linux,revision appears to be in big endian
    #endif
    
    for (size_t i = 0; i < rpi_hw_info_size; ++i) {
        if (rev == rpi_hw_info[i].hwver) {
            result = &rpi_hw_info[i];
            break;
        }
    }
#else
    file.open("/proc/cpuinfo");
    if (!file.is_open()) {
        return nullptr;
    }
    
    while (std::getline(file, line)) {
        if (line.find(HW_VER_STRING) != std::string::npos) {
            std::string::size_type pos = line.find(": ");
            if (pos == std::string::npos) {
                continue;
            }
            
            std::string rev_str = line.substr(pos + 2);
            std::stringstream ss;
            ss << std::hex << rev_str;
            ss >> rev;
            
            if (ss.fail()) {
                continue;
            }
            
            for (size_t i = 0; i < rpi_hw_info_size; ++i) {
                uint32_t hwver = rpi_hw_info[i].hwver;
                
                // Take out warranty and manufacturer bits
                hwver &= ~(RPI_WARRANTY_MASK | RPI_MANUFACTURER_MASK);
                rev &= ~(RPI_WARRANTY_MASK | RPI_MANUFACTURER_MASK);
                
                if (rev == hwver) {
                    result = &rpi_hw_info[i];
                    break;
                }
            }
        }
    }
#endif
    file.close();
    return result;
}
